#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Feb 15 20:28:22 2021

@author: liqingsimac
"""

#########################################
## 8.2. 基本思想  
'''
import numpy as np
import matplotlib.pyplot as plt

def myfun(x,lam):
    y0=0
    return (y0-1)*np.exp(lam*x)+np.exp(-x)

x=np.linspace(0,1,31)
lam1=-5
lam2=-10
lam3=-30
lam4=-200
y1=myfun(x,lam1)
y2=myfun(x,lam2)
y3=myfun(x,lam3)
y4=myfun(x,lam4)
fig=plt.figure()
ax=fig.add_subplot(111)
ax.plot(x,y1,'-',label='lambda=-5')
ax.plot(x,y2,'-.',label='lambda=-10')
ax.plot(x,y3,'--',label='lambda=-30')
ax.plot(x,y4,':',label='lambda=-200')
ax.legend()
'''

#########################################
## 8.3.2. 谐波振荡器ex.1 
'''
import numpy as np
from scipy.integrate import odeint
f=lambda y,t:y
odeint(f,1,[0,1])
'''

#########################################
##8.3.2. 谐波振荡器ex.2
'''
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

def rhs(Y,t,omega):
    y,ydot=Y
    return ydot,-omega**2*y

t_arr=np.linspace(0,2*np.pi,101)
y_init=[1,0]
omega=2.0
y_arr=odeint(rhs,y_init,t_arr,args=(omega,))
y,ydot=y_arr[:,0],y_arr[:,1]

fig=plt.figure()
ax1=fig.add_subplot(121)
ax1.plot(t_arr,y,t_arr,ydot)
ax1.set_xlabel('t')
ax1.set_ylabel('y and ydot')
ax2=fig.add_subplot(122)
ax2.plot(y,ydot)
ax2.set_xlabel('y')
ax2.set_ylabel('ydot')
plt.suptitle('Solution curve when omega =%2g' % omega)
fig.tight_layout()
fig.subplots_adjust(top=0.90)
'''

#########################################
##8.3.2.ex.2.ver.2
'''
import numpy as np
import matplotlib.pyplot as plt 
from scipy.integrate import odeint

def rhs(Y, t, omega):
    y,ydot=Y
    return ydot, -omega**2*y

t_arr=np.linspace(0, 2*np.pi, 101)
y_init=[1, 0]
omega=2.0

fig=plt.figure()
ax=fig.add_subplot(111)
y,ydot=np.mgrid[-3:3:21j, -6:6:21j]
u,v=rhs(np.array([y, ydot]), 0.0, omega)
mag=np.hypot(u, v)
mag[mag==0]=1.0
ax.quiver(y, ydot, u/mag, v/mag, color='blue')

# Enable drawing of arbitrary number of trajectories 
print('\n\n\nUse mouse to select each starting point')
print('Timeout after 30 seconds')
choice=[(0,0)]
while len(choice) > 0:
    y01 = np.array([choice[0][0], choice[0][1]]) 
    y = odeint(rhs, y01, t_arr, args=(omega,)) 
    plt.plot(y[:, 0], y[:, 1], lw=2)
    choice = plt.ginput()
print('Timed out!')
'''

#########################################
## 8.3.3. van de pol equation
'''
import numpy as np
from scipy.integrate import odeint

def rhs(z,t,mu):
    return [ z[1], mu*(1-z[0]**2)*z[1]-z[0] ]

def jac(z,t,mu):
    return[ [0,1], [-2*mu*z[0]*z[1]-1,mu*(1-z[0]**2)] ]

mu=1.0
t_final = 15.0 if mu<10 else 4.0*mu
n_points = 1001 if mu<10 else 1001*mu
t=np.linspace(0,t_final,n_points)
y0=np.array([0.1,0.0])
y,info=odeint(rhs,y0,t,args=(mu,),Dfun=jac,full_output=True)

print('mu=%g, number of Jacobian calls is %d' %
      (mu,info['nje'][-1]) )

import matplotlib.pyplot as plt
plt.plot(y[:,0],y[:,1])
'''

#########################################
##8.3.4. Lorenz equation
'''
import numpy as np
from scipy.integrate import odeint

def rhs(u,t,beta,rho,sigma):
    x,y,z=u
    return [sigma*(y-x), rho*x-y-x*z, x*y-beta*z ]

sigma=10.0
beta=8.0/3.0
rho1=29.0
rho2=28.8

u01=[1.0,1.0,1.0]
u02=[1.0,1.0,1.0]

t=np.linspace(0.0, 50.0, 10001)
u1=odeint(rhs,u01,t,args=(beta,rho1,sigma))
u2=odeint(rhs,u02,t,args=(beta,rho2,sigma))

x1,y1,z1=u1[:,0],u1[:,1],u1[:,2]
x2,y2,z2=u2[:,0],u2[:,1],u2[:,2]

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig=plt.figure()
ax=Axes3D(fig)
ax.plot(x1,y1,z1,'b-')
#ax.plot(x2,y2,z2,'g:')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
#ax.set_title('Lorenz equations with rho=%g and rho=%g'%(rho1,rho2))
ax.set_title('Lorenz equations with rho=%g'%rho1)
'''


#########################################
##8.4.3. a simple example 两点边值问题的简单例子：简谐振动
'''
import numpy as np
import scikits.bvp1lg.colnew as colnew

degrees=[2]
boundary_points=np.array([0,np.pi])
tol=1.0e-8*np.ones_like(boundary_points)

def fsub(x,Z):
    #The equations
    u,du=Z
    return np.array([-u])

def gsub(Z):
    #The boundary conditions
    u,du=Z
    return np.array([u[0],du[1]-1.0])

solution=colnew.solve(boundary_points,degrees,fsub,gsub,
                      is_linear=True,tolerances=tol,
                      vectorized=True,maximum_mesh_size=300)
x=solution.mesh
u_exact=-np.sin(x)

import matplotlib.pyplot as plt
fig=plt.figure()
ax=fig.add_subplot(111)
ax.plot(x,solution(x)[:,0],'b.')
ax.plot(x,u_exact,'g-')
'''


#########################################
## 8.4.3. example from the web, Bessel equation.
'''
import numpy as np
import scipy.special as special
import scikits.bvp1lg.colnew as colnew

nu = 3.4123
degrees = [2, 1]
def fsub(x, z):
    u, du, v = z     # it's neat to name the variables
    return np.array([-du/x + (nu**2/x**2 - 1)*u, 
                     x**(nu+1) * u])

def dfsub(x, z):
    u, du, v = z
    zero = np.zeros(x.shape)
    return np.array([[(nu**2/x**2 - 1), -1/x, zero],
                      [       x**(nu+1), zero, zero]])

boundary_points = [1, 5, 10]

def gsub(z):
    u, du, v = z
    return np.array([u[0] - special.jv(nu,   1),
            v[1] - 5**(nu+1) * special.jv(nu+1, 5),
            u[2] - special.jv(nu,   10)])

def dgsub(z):
    return np.array([[1, 0, 0],
                     [0, 0, 1],
                     [1, 0, 0]])

tol = [1e-5, 0, 1e-5]
solution = colnew.solve(
        boundary_points, degrees, fsub, gsub,
        dfsub=dfsub, dgsub=dgsub,
        is_linear=True, tolerances=tol,
        vectorized=True, maximum_mesh_size=300)

solution.nmesh

x = np.linspace(1, 10, 101)
np.allclose(solution(x)[:,0], special.jv(nu, x),
            rtol=1e-4, atol=1e-8)

np.allclose(solution(x)[:,2], 
            x**(nu+1)*special.jv(nu+1, x),
            rtol=1e-4, atol=1e-8)
            
import matplotlib.pyplot as plt
plt.plot(solution.mesh, 
         solution(solution.mesh)[:,2], '.',
         x, x**(nu+1)*special.jv(nu+1, x), '-')
         
#plt.show()
'''

#########################################
## 8.4.4. 线性特征值问题 
'''
import numpy as np
import scikits.bvp1lg.colnew as colnew
degrees=[2,1,1]
boundary_points=np.array([0.0, 0.0, 1.0, 1.0])
tol=1.0e-5*np.ones_like(boundary_points)

def fsub(x,Z):
    u,du,v,w=Z
    ddu=-u*v
    dv=np.zeros_like(x)
    dw=u*u
    return np.array([ddu,dv,dw])

def gsub(Z):
    u,du,v,w=Z
    return np.array([u[0],w[1],u[2],w[3]-1.0])

guess_lambda=100.0 
def guess(x):
    u=x*(1.0/3.0-x)*(2.0/3.0-x)*(1.0-x)
    du=2.0*(1.0-2.0*x)*(1.0-9.0*x+9.0*x*x)/9.0
    v=guess_lambda*np.ones_like(x)
    w=u*u
    Z_guess=np.array([u,du,v,w]) 
    f_guess=fsub(x,Z_guess) 
    return Z_guess,f_guess

solution=colnew.solve(boundary_points,degrees,fsub,gsub,
                   is_linear=False,initial_guess=guess,
                   tolerances=tol,vectorized=True,
                   maximum_mesh_size=300)

# plot solution

import matplotlib.pyplot as plt 
#%matplotlib notebook

#plt.ion()
x=solution.mesh
u_exact=np.sqrt(2)*np.sin(3.0*np.pi*x) 
plt.plot(x,solution(x)[:,0],'b.',x,u_exact,'g-')
print('Third eigenvalue is %16.10e.'%solution(x)[0,2])

'''

#########################################
##8.4.5. Bratu problem
'''
import numpy as np
import scikits.bvp1lg.colnew as colnew

degree=[2]
boundary_points=np.array([0.0, 1.0])
tol=1.0e-8*np.ones_like(boundary_points)

lamda=1.0
mu=0.2

def fsub(x,Z):
    u,du=Z
    ddu=-lamda*np.exp(u)
    return np.array([ddu])

def gsub(Z):
    u,du=Z
    return np.array([u[0],u[1]])

def initial_guess(x):
    u=mu*x*(1.0-x)
    du=mu*(1.0-2.0*x)
    Z_guess=np.array([u,du])
    f_guess=fsub(x,Z_guess)
    return Z_guess,f_guess

solution=initial_guess

solution=colnew.solve(boundary_points,degrees,fsub,gsub,
                      is_linear=False,initial_guess=solution,
                      tolerances=tol,vectorized=True,
                      maximum_mesh_size=300)
    
import matplotlib.pyplot as plt
#%matplotlib notebook

plt.ion()
x=solution.mesh
plt.plot(x,solution(x)[:,0],'b.')
plt.show(block=False)
'''

#########################################
##8.4.5. Bratu方程
'''
import numpy as np
import scikits.bvp1lg.colnew as colnew

degrees=[2,1,1]
boundary_points=np.array([0.0,0.0,1.0,1.0])
tol=1.0e-8*np.ones_like(boundary_points)

def fsub(x,Z):
    u,du,v,w=Z
    ddu=-v*np.exp(u)
    dv=np.zeros_like(x)
    dw=u*u
    return np.array([ddu,dv,dw])

def gsub(Z):
    u,du,v,w=Z
    return np.array([u[0],w[1],u[2],w[3]-gamma])


def guess(x):
    u=0.5*x*(1.0-x) 
    du=0.5*(1.0-2.0*x) 
    v=np.zeros_like(x)
    w=u*u 
    Z_guess=np.array([u,du,v,w]) 
    f_guess=fsub(x,Z_guess) 
    return Z_guess,f_guess

solution=guess
gaml=[]
laml=[]

for gamma in np.linspace(0.01,5.01,1001):
    solution = colnew.solve(boundary_points,degrees,fsub,gsub,
                        is_linear=False,initial_guess=solution,
                        tolerances=tol,vectorized=True,
                        maximum_mesh_size=300)
    x=solution.mesh 
    lam=solution(x)[:,2] 
    gaml.append(gamma) 
    laml.append(np.max(lam))

# plot solution
import matplotlib.pyplot as plt 
#%matplotlib notebook
plt.ion()
fig = plt.figure()
ax = fig.add_subplot(111) 
ax.plot(gaml, laml) 
ax.set_xlabel(r'$\gamma$', size=20)
ax.set_ylabel(r'$\lambda$', size=20) 
ax.grid(b=True)                 
'''

#########################################
## 8.4.5. Bratu 方程的另一种数值求解，使用 
## scipy.integrate.solve_bvp 函数
''' 
import numpy as np

def fun(x, y):
    return np.vstack((y[1], -np.exp(y[0])))

def bc(ya, yb):
    return np.array([ya[0], yb[0]])

x = np.linspace(0, 1, 5)

y_a = np.zeros((2, x.size))
y_b = np.zeros((2, x.size))
y_b[0] = 3

from scipy.integrate import solve_bvp
res_a = solve_bvp(fun, bc, x, y_a)
res_b = solve_bvp(fun, bc, x, y_b)

x_plot = np.linspace(0, 1, 100)
y_plot_a = res_a.sol(x_plot)[0]
y_plot_b = res_b.sol(x_plot)[0]

import matplotlib.pyplot as plt
plt.plot(x_plot, y_plot_a, label='y_a')
plt.plot(x_plot, y_plot_b, label='y_b')
plt.legend()
plt.xlabel('x')
plt.ylabel('y')
plt.show()
'''



#########################################
##8.5.3. delay logistic equation
'''
import numpy as np
import matplotlib.pyplot as plt 
from pydelay import dde23

t_final=50
delay=2.0
x_initial=0.1
equations={'x': 'x*(1.0-x(t-tau))'}
parameters={'tau' : delay}
dde=dde23(eqns=equations,params=parameters)

dde.set_sim_params(tfinal=t_final,dtmax=1.0,
                   AbsTol=1.0e-6,RelTol=1.0e-3)
histfunc={'x': lambda t: x_initial}
dde.hist_from_funcs(histfunc,101)
dde.run()

t_vis=0.1*t_final
sol=dde.sample(tstart=t_vis+delay,tfinal=t_final,dt=0.1)
t=sol['t']
x=sol['x']
sold=dde.sample(tstart=t_vis,tfinal=t_final-delay,dt=0.1)
xd=sold['x']

#plt.ion()
fig = plt.figure()
ax1=fig.add_subplot(121)
ax1.plot(t, x)
ax1.set_xlabel('t')
ax1.set_ylabel('x(t)')

ax2=fig.add_subplot(122)
ax2.plot(x, xd)
ax2.set_xlabel('tx')
ax2.set_ylabel('x(t-tau)')
'''

######################################################
## Downloaded April 11, 2022
## http://pydelay.sourceforge.net
#'''
# import pydelay and numpy and pylab
import numpy as np
import pylab as pl
from pydelay import dde23

# define the equations
eqns = {
    'x' : '0.25 * x(t-tau) / (1.0 + pow(x(t-tau),p)) -0.1*x'
    }

#define the parameters
params = {
    'tau': 15,
    'p'  : 10
    }

# Initialise the solver
dde = dde23(eqns=eqns, params=params)

# set the simulation parameters
# (solve from t=0 to t=1000 and limit the maximum step size to 1.0)
dde.set_sim_params(tfinal=1000, dtmax=1.0)

# set the history of to the constant function 0.5 (using a python lambda function)
histfunc = {
    'x': lambda t: 0.5
    }
dde.hist_from_funcs(histfunc, 51)

# run the simulator
dde.run()

# Make a plot of x(t) vs x(t-tau):
# Sample the solution twice with a stepsize of dt=0.1:

# once in the interval [515, 1000]
sol1 = dde.sample(515, 1000, 0.1)
x1 = sol1['x']

# and once between [500, 1000-15]
sol2 = dde.sample(500, 1000-15, 0.1)
x2 = sol2['x']

pl.plot(x1, x2)
pl.xlabel('$x(t)$')
pl.ylabel('$x(t - 15)$')
pl.show()
#'''



################################################
##8.6.1. wiener process
'''
import numpy as np
import numpy.random as npr
T=1
N=500
t,dt=np.linspace(0,T,N+1,retstep=True)
dW=npr.normal(0.0,np.sqrt(dt),N+1)
dW[0]=0.0
W=np.cumsum(dW)
import matplotlib.pyplot as plt 
#%matplotlib notebook
#plt.ion()
plt.plot(t,W)
plt.xlabel('t')
plt.ylabel('W(t)')
plt.title('Sample Wiener Process',weight='bold',size=12)
'''

#########################################
##8.6.2. ito公式的一个简单例子，均值函数的解析解与模拟解
'''
import numpy as np
import numpy.random as npr
T=1
N=1000
M=5000
t,dt=np.linspace(0,T,N+1,retstep=True)
dW=npr.normal(0.0,np.sqrt(dt),(M,N+1))
dW[ : ,0]=0.0
W=np.cumsum(dW,axis=1)
U=np.exp(- 0.5*W)
Umean=np.mean(U,axis=0)
Uexact=np.exp(t/8)

import matplotlib.pyplot as plt 
#%matplotlib notebook
#plt.ion()

fig=plt.figure()
ax=fig.add_subplot(111)
ax.plot(t,Umean,'b-',label='mean of %d paths' % M)
ax.plot(t,Uexact,'r-.', label='exact ' + r'$\langle U\rangle$')
for i in range(5):
    ax.plot(t,U[i,:],'g--')

ax.set_xlabel('t')
ax.set_ylabel('U')
ax.legend(loc='best')

maxerr=np.max(np.abs(Umean-Uexact))
print('With %d paths and %d intervals the max error is %g' % \
      (M,N,maxerr))

#'''

#########################################
##8.6.4.a BS方程得数值解，Milstein方法
'''
import numpy as np
import numpy.random as npr

#set up grid
T=1.0
N=30
t,dt=np.linspace(0,T,N+1,retstep=True)

#Get Brownian motion
dW=npr.normal(0.0,np.sqrt(dt),N+1)
dW[0]=0.0
W=np.cumsum(dW)

#Equation parameters and functions
lamda=2.0
mu=1.0
Xzero=1.0
def a(X): return lamda*X
def b(X): return mu*X
def bd(X): return mu*np.ones_like(X)

#Analytic solution
Xanal=Xzero*np.exp((lamda-0.5*mu*mu)*t+mu*W)

#Milstein solution
Xmil=np.empty_like(t)
Xmil[0]=Xzero
for n in range(N):
    Xmil[n+1]=Xmil[n] + dt*a(Xmil[n]) + dW[n+1]*b(Xmil[n]) \
    +0.5*(b(Xmil[n])*bd(Xmil[n])*(dW[n+1]**2-dt))

import matplotlib.pyplot as plt
plt.ion()
plt.plot(t,Xanal,'b-',label='analytic')
plt.plot(t,Xmil,'g-.',label='Milstein')
plt.legend(loc='best')
plt.xlabel('t')
plt.ylabel('X(t)')
'''

#########################################
##8.6.4.b
'''
import numpy as np
import numpy.random as npr

#Problem definition
M=1000  #number of paths sampled
P=6  #number of discretizations
T=1  #endpoint of time interval
N=2**12  #finest grid size
dt=1.0*T/N

#Problem parameters
lamda=2.0
mu=1.0
Xzero=1.0

def a(X): return lamda*X
def b(X): return mu*X
def bd(X): return mu*np.ones_like(X)

#build the Brownian paths
dW=npr.normal(0.0, np.sqrt(dt),(M,N+1))
dW[:,0]=0.0
W=np.cumsum(dW,axis=1)

#build the exact solutions at the end of the paths
ones=np.ones(M)
Xexact=Xzero*np.exp((lamda-0.5*mu*mu)*ones+mu*W[:,-1])
Xemerr=np.empty((M,P))
Xmilerr=np.empty((M,P))

#loop over refinements
for p in range(P):
    R=2**p
#    L=int(N/R)
    L=2**(12-p);Dt=R*dt
    Xem=Xzero*ones;Xmil=Xzero*ones
    Wc=W[:,::R]
    for j in range(L):
        deltaW=Wc[:,j+1]-Wc[:,j]
        Xem+=Dt*a(Xem)+deltaW*b(Xem)
        Xmil+=Dt*a(Xmil)+deltaW*b(Xmil)+0.5*b(Xmil)*bd(Xmil)*(deltaW**2-Dt)
    Xemerr[:,p]=np.abs(Xem-Xexact)
    Xmilerr[:,p]=np.abs(Xmil-Xexact)
    
#do some plotting
import matplotlib.pyplot as plt
Dtvals=dt*np.array([2**p for p in range(P)])
lDtvals=np.log10(Dtvals)
Xemerrmean=np.mean(Xemerr,axis=0)
plt.plot(lDtvals,np.log10(Xemerrmean),'bo')
plt.plot(lDtvals,np.log10(Xemerrmean),'b:',label='EM actual')
plt.plot(lDtvals,0.5*np.log10(Dtvals),'b-.',label='EM theoretical') 
       
Xmilerrmean=np.mean(Xmilerr,axis=0)
plt.plot(lDtvals,np.log10(Xmilerrmean),'go')
plt.plot(lDtvals,np.log10(Xmilerrmean),'g:',label='Mil actual')
plt.plot(lDtvals,np.log10(Dtvals),'g-.',label='Mil theoretical') 
plt.legend(loc='best')
#plt.xlabel('log of Delta t',size=12)
#plt.ylabel('log of expectation of difference',size=12)
plt.xlabel(r'$\log_{10}\Delta t$',size=16) 
plt.ylabel(r'$\log_{10}\left(\langle|X_n-X(\tau)|\rangle\right)$',size=16)

emslope=((np.log10(Xemerrmean[-1]) - np.log10(Xemerrmean[0]))/
         (lDtvals[-1]-lDtvals[0]))
print('Emprical EM slope is %g' %emslope)
milslope=((np.log10(Xmilerrmean[-1]) - np.log10(Xmilerrmean[0]))/
          (lDtvals[-1]-lDtvals[0]))
print('Emprical mil slope is %g' %milslope)
'''









